package com.hccl.service.classifier;

/**
 * Created by yang on 2018/8/3.
 */
import com.hccl.service.classifier.ClassifierResult;
import org.tensorflow.Graph;
import org.tensorflow.Session;
import org.tensorflow.Tensor;
import org.tensorflow.Tensors;
import org.tensorflow.framework.ConfigProto;

import java.io.IOException;
import java.nio.file.Files;
import java.nio.file.Path;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;

import static com.hccl.config.Constants.sentlen;

public class runRnnSession extends runSession{
    Graph g;
    ConfigProto config;

    public runRnnSession(Path rnnModelPath) throws IOException {

        g = new Graph();

        g.importGraphDef(Files.readAllBytes(rnnModelPath));
        //System.out.print("read pb model done");
        config = ConfigProto.newBuilder().setAllowSoftPlacement(true).build();
        System.out.print("read pb model done");

    }

    public List<ClassifierResult> run(Map featureMap) throws Exception {

        int[][] input_word = new int[1][15];
        int[][] input_pos = new int[1][15];
        int[][] input_position = new int[1][15];
        //long[][] input_sentlen = new long[1][doclen];
        //long[] dropout = new double[1];
        int[][] wordData = new int[1][sentlen];
        int[][] posData = new int[1][sentlen];
        int[][] positionData = new int[1][sentlen];
        wordData= (int[][])featureMap.get("data");
        posData= (int[][])featureMap.get("posdata");
        positionData= (int[][])featureMap.get("positiondata");
        input_word[0] = wordData[0];
        input_pos[0] = posData[0];
        input_position[0] = positionData[0];
        float drop_out = (float)1.0;
        //input_sentlen[0] = (long[]) featureMap.get("sentlen");

        List<ClassifierResult> result = new ArrayList<>();

        try(Tensor inputsWord = Tensors.create(input_word);
            Tensor inputsPos = Tensors.create(input_pos);
            Tensor inputsPosition = Tensors.create(input_position);
            Tensor dropout = Tensors.create(drop_out);
            //Tensor inputsSentlen = Tensors.create(input_sentlen);
            Session session = new Session(g,config.toByteArray())) {
            Session.Runner run = session.runner().feed("input_x:0", inputsWord)
                                                .feed("input_pos:0", inputsPos)
                                                .feed("input_position:0", inputsPosition)
                                                .feed("dropout_keep_prob:0", dropout);
            result = getFetch(run);
        }

        return result;

    }
}
